# -*- coding: utf-8 -*-
'''
Created on Sep 1, 2019

@author: yl
'''
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.ticker import MaxNLocator
from FFT_Interpolation import FFT_interpolation_2, line_cal, FFT_interpolation_boxcar
from scipy.optimize import curve_fit
from Frames_set_Compare import Frame_set, frame_total, fs_cam, V_r_x, V_r_y, V_r_z, j, c, Lamda, Fc, k, I_0, pix_size, pix_num, screen_diameter

def fit_func(x, a, b, c):
    return a*(x-b)**2 + c
# p= -0.0213125
p= -0.02135
def fit_func_phase(x, q):
    return p*x+q

def fit_func_line(x, p, q):
    return p*x+q


img_set = Frame_set
V_x, V_y, V_z = V_r_x, V_r_y, V_r_z

frame_num = 0
hor_threshold, ver_threshold = 30, 30
hor_centerline, ver_centerline = img_set[int(frame_num/2)][int(len(img_set[0])/2), :], img_set[int(frame_num/2)][:, int(len(img_set[0][1])/2)]
hor_start, ver_start = np.where(hor_centerline > hor_threshold)[0][0], np.where(ver_centerline > ver_threshold)[0][0]
hor_end, ver_end = np.where(hor_centerline > hor_threshold)[0][-1], np.where(ver_centerline > ver_threshold)[0][-1]
hor_index = int((ver_start+ver_end)/2)
ver_index = int((hor_start+hor_end)/2)

line_num = 51
scale_factor = 4
hor_lines = np.linspace(hor_index-scale_factor*(line_num-1)/2, hor_index+scale_factor*(line_num-1)/2, line_num, dtype='int')
# hor_lines = np.linspace(470, 500, 31, dtype='int')
ver_lines = np.linspace(ver_index-scale_factor*(line_num-1)/2, ver_index+scale_factor*(line_num-1)/2, line_num, dtype='int')
# print(hor_lines)
# print(ver_lines)

hor_f_set, ver_f_set = [], []
hor_phi_set, ver_phi_set = [], []

for img in img_set:
# for img in img_set[0:2]:
    frame_num += 1
    for i in range(line_num):
        hor_center = np.array(img[hor_lines[i]])
        ver_center = img[:, ver_lines[i]]
        
        hor_f_fit, hor_phase_estim, hor_m_k_num = line_cal(hor_center)
        ver_f_fit, ver_phase_estim, ver_m_k_num = line_cal(ver_center)

        hor_f_set.append(hor_f_fit)
        hor_phi_set.append(hor_phase_estim)
        ver_f_set.append(ver_f_fit)
        ver_phi_set.append(ver_phase_estim)
        
        print(frame_num, hor_m_k_num, ver_m_k_num)
hor_f_set = np.array(hor_f_set)
ver_f_set = np.array(ver_f_set)
hor_f_set = hor_f_set.reshape(line_num, frame_num, order='F')
ver_f_set = ver_f_set.reshape(line_num, frame_num, order='F')
hor_angle_set = (V_x-Lamda*hor_f_set/2)*1e6 ### urad
ver_angle_set = (V_x-Lamda*ver_f_set/2)*1e6 ### urad

hor_non_f_set, hor_non_phi_set = [], []
ver_non_f_set, ver_non_phi_set = [], []
for i in range(line_num):
    hor_non_f, hor_non_phi = FFT_interpolation_boxcar(hor_angle_set[i]-np.average(hor_angle_set[i]), 1/frame_num)[0:2]
    ver_non_f, ver_non_phi = FFT_interpolation_boxcar(ver_angle_set[i]-np.average(ver_angle_set[i]), 1/frame_num)[0:2]
    hor_non_f_set.append(hor_non_f)
    hor_non_phi_set.append(hor_non_phi)
    ver_non_f_set.append(ver_non_f)
    ver_non_phi_set.append(ver_non_phi)

hor_amp = []
ver_amp = []
for i in range(line_num):
    hor_amp_line = np.max(hor_angle_set[i])-np.min(hor_angle_set[i])
    ver_amp_line = np.max(ver_angle_set[i])-np.min(ver_angle_set[i])
    hor_amp.append(hor_amp_line)
    ver_amp.append(ver_amp_line)
 
hor_amp = np.array(hor_amp)
ver_amp = np.array(ver_amp)
hor_used_line = hor_lines[np.where(hor_amp < 3)[0]]
ver_used_line = ver_lines[np.where(ver_amp < 3)[0]]
# print (repr(hor_used_line))
# print (repr(ver_used_line))


y = [np.arange(frame_num)]*line_num
y = np.array(y)

fig = plt.figure('hor_tilt')
x = [hor_lines]*frame_num
x = np.array(x).T
ax = fig.gca(projection='3d')
plt.gca().patch.set_facecolor('white')
plt.title('hor_tilt')
for i in range(line_num):
    ax.plot(x[i,:],y[i,:],hor_angle_set[i,:])
ax.xaxis.set_major_locator(MaxNLocator(integer=True))

fig = plt.figure('ver_tilt')
x = [ver_lines]*frame_num
x = np.array(x).T
ax = fig.gca(projection='3d')
plt.gca().patch.set_facecolor('white')
plt.title('ver_tilt')
plt.xlabel('line number')
plt.ylabel('frame number')
for i in range(line_num):
    ax.plot(x[i,:],y[i,:],ver_angle_set[i,:])
ax.xaxis.set_major_locator(MaxNLocator(integer=True))

# plt.figure('nonlinearity fft')
# plt.subplot(2,2,1)
# plt.plot(hor_lines, hor_non_f_set)
# plt.grid(which='both', axis='both')
# plt.subplot(2,2,3)
# plt.plot(hor_lines, (hor_non_phi_set))
# plt.grid(which='both', axis='both')
# plt.subplot(2,2,2)
# plt.plot(ver_lines, ver_non_f_set)
# plt.grid(which='both', axis='both')
# plt.subplot(2,2,4)
# plt.plot(ver_lines, (ver_non_phi_set))
# plt.grid(which='both', axis='both')

plt.figure('nonlinearity amp')
ax1 = plt.subplot(2,1,1)
ax1.plot(hor_lines, hor_amp)
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.title('hor tilt nonlinearity amplitude')
plt.grid(which='both', axis='both')
plt.xlabel('line number')
plt.ylabel('amplitude(urad)')
ax2 = plt.subplot(2,1,2)
plt.title('ver tilt nonlinearity amplitude')
plt.plot(ver_lines, ver_amp)
ax2.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.grid(which='both', axis='both')
plt.xlabel('line number')
plt.ylabel('amplitude(urad)')
plt.tight_layout()
plt.show()